import torch
import torch.nn as nn


class ChannelWiseAttention(nn.Module):
    def __init__(self):
        super(ChannelWiseAttention, self).__init__()
        self.attention_layer = nn.Conv1d(1, 1, kernel_size=1)

    def forward(self, x):
        # batch_size, num_samples, num_channels = x.size()
        x = x.permute(0, 2, 1).unsqueeze(1)  # Reshape to (batch_size, 1, num_channels, num_samples)

        attention_scores = self.attention_layer(x)
        attention_weights = torch.softmax(attention_scores, dim=2)

        attended_input = x * attention_weights

        output = attended_input + x  # Residual connection
        output = output.squeeze(1).permute(0, 2, 1)  # Reshape back to (batch_size, num_samples, num_channels)

        return output


class EEGEncoding(nn.Module):
    def __init__(self, num_channels, hidden_size=64, num_layers=2, num_classes=2):
        super(EEGEncoding, self).__init__()
        self.channel_wise_attention = ChannelWiseAttention()
        self.conv1 = nn.Conv1d(num_channels, hidden_size, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(hidden_size, hidden_size, kernel_size=3, padding=1)
        self.rnn = nn.GRU(hidden_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_size * 2, num_classes)

    def forward(self, x):
        batch_size, num_samples, num_channels = x.size()
        x = self.channel_wise_attention(x)
        x = x.permute(0, 2, 1)  # Reshape to (batch_size, num_channels, num_samples)
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = x.permute(0, 2, 1)  # Reshape to (batch_size, num_samples, hidden_size)
        _, hidden = self.rnn(x)
        hidden = hidden.permute(1, 0, 2).contiguous().view(batch_size, -1)
        x = self.fc(hidden)
        return x
